import java.util.Arrays;

public class Solution {
    public static void main(String[] args) {
        Solution s = new Solution();
        int[] vals = {1,3,2,1,3};
        int[][] edges = {
                {0, 1},
                {0, 2},
                {2, 3},
                {2, 4}
        };
        System.out.println(s.numberOfGoodPaths(vals, edges));
    }

    public int numberOfGoodPaths(int[] vals, int[][] edges) {
        /**
         * 好路径的数目*/
        // 1 预处理（按照值从小到大排升序）
        int n = vals.length;
        initialize(n, vals);
        // --! 错误在此处：原本的逻辑是按照值从小达到排列，但是我现在实际是按节点的序号进行排列，正确应当是比较对应的vals!--
        Arrays.sort(edges, (a,b) -> Math.max(vals[a[0]], vals[a[1]]) - Math.max(vals[b[0]], vals[b[1]]));

        // 2 从低到高合并集合
        int answer = 0;
        for(int[] e: edges) {
            int fa = find(e[0]);
            int fb = find(e[1]);
            // -2.1 若是集合最大值相同，则计算好路径
            if(maxValue[fa] == maxValue[fb]) {
                answer = answer + (maxValueCount[fa] * maxValueCount[fb]);
            }
            // -2.2 合并连接
            union(e[0], e[1]);
        }

        // 3 返回值（单独加上节点个数）
        return answer + n;
    }

    // 创建简易并查集
    public static int MAX_N;
    public static int[] father;
    public static int[] size;
    public static int[] stack;
    // -标记集合最大值与最大值个数
    public static int[] maxValue;
    public static int[] maxValueCount;

    public static void initialize(int n, int[] vals) {
        /**
         * 初始化并查集*/
        MAX_N = n;
        father = new int[MAX_N];
        size = new int[MAX_N];
        stack = new int[MAX_N];
        maxValue = new int[MAX_N];
        maxValueCount = new int[MAX_N];
        for(int i = 0; i < MAX_N; i++) {
            father[i] = i;
            size[i] = 1;
            maxValue[i] = vals[i];
            maxValueCount[i] = 1;
        }
    }

    public static int find(int a) {
        /**
         * 找到a所在集合的代表元素+扁平化*/
        int i = 0;
        while(a != father[a]) {
            stack[i++] = a;
            a = father[a];
        }
        while(i > 0) {
            father[stack[--i]] = a;
        }
        return a;
    }

    public static boolean isSameSet(int a, int b) {
        /**
         * 判断a、b元素是否属于同一集合*/
        return find(a) == find(b);
    }

    public static void union(int a, int b) {
        /**
         * 合并a、b元素所在集合+小挂大*/
        int fa = find(a);
        int fb = find(b);
        if(fa != fb) {
            if(size[fa] >= size[fb]) {
                size[fa] += size[fb];
                // -更新最大值与最大值个数
                if(maxValue[fa] == maxValue[fb]) {
                    maxValueCount[fa] += maxValueCount[fb];
                } else if(maxValue[fa] < maxValue[fb]) {
                    maxValue[fa] = maxValue[fb];
                    maxValueCount[fa] = maxValueCount[fb];
                }
                father[fb] = fa;
            } else {
                size[fb] += size[fa];
                // -更新最大值与最大值个数
                if(maxValue[fa] == maxValue[fb]) {
                    maxValueCount[fb] += maxValueCount[fa];
                } else if(maxValue[fb] < maxValue[fa]) {
                    maxValue[fb] = maxValue[fa];
                    maxValueCount[fb] = maxValueCount[fa];
                }
                father[fa] = fb;
            }
        }
    }
}
